package com.starhub.application.agent.rag;

import java.util.Hashtable;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import com.starhub.application.rag.dto.MetaRagDto;
import com.starhub.application.rag.dto.RetrieveDto;
import com.starhub.application.rag.enums.RetrieveEnums;
import com.starhub.application.rag.result.RagResult;
import com.starhub.application.rag.service.RagService;
import com.starhub.application.rag.util.EmbeddingSearchResultUtil;
import com.starhub.common.bean.agent.PromptVectorConfig;
import com.starhub.utils.Util;

import dev.langchain4j.data.segment.TextSegment;
import lombok.extern.slf4j.Slf4j;

/**
 * 智能体内的召回
 * 
 * @version: V1.0
 * @author: liuhf
 * @time 2025-07-09
 */

@Slf4j
@Component
public class RagHelper {

    @Autowired
    private RagService ragService;

    public Map<String, String> search(String queryText, List<PromptVectorConfig> vectorConfigList) {
        Map<String, String> result = new Hashtable<String, String>();

        //召回内部存在多线程处理，这里暂时不使用多线程，循环进行召回，后续需要优化（MQ）
        vectorConfigList.stream().forEach(vectorConfig -> {
            try{
                List<TextSegment> list = null;
                if (vectorConfig.getRetrieveType() == RetrieveEnums.EMBEDDING.getCode()
                        || vectorConfig.getRetrieveType() == RetrieveEnums.AGENT.getCode()) {
                    // 增强召回模式
                    MetaRagDto metaRagDto = MetaRagDto.builder()
                            .mainCollectionId(Util.null2String(vectorConfig.getKnowledgeCollectionId()))
                            .retrieveType(RetrieveEnums.getRetrieveEnums(vectorConfig.getRetrieveType()))
                            .queryText(queryText)
                            .build();
                    list = ragService.retrieveByMetaData(ragService.checkMetaRagDto(metaRagDto));
                } else {
                    // 普通召回模式
                    RetrieveDto retrieveDto = RetrieveDto.builder()
                            .queryText(queryText)
                            .topK(vectorConfig.getMaxResults())
                            .minScore(vectorConfig.getSimilary())
                            .collectionId(Util.null2String(vectorConfig.getKnowledgeCollectionId()))
                            .build();
                    RagResult ragResult = ragService.retrieve(ragService.checkRetrieveDto(retrieveDto));
                    list = EmbeddingSearchResultUtil.transToSegments(ragResult.getSearchResult());
                }
                // 获取召回结果
                String resultString = getResult(list);
                result.put(Util.null2String(vectorConfig.getRetrieveKey()), resultString);

            }catch(Exception e){
                log.error("智能体RAG召回失败 - CollectionId: {}, RetrieveType: {}, QueryText: {}, Error: {}", 
                    vectorConfig.getKnowledgeCollectionId(), 
                    vectorConfig.getRetrieveType(), 
                    queryText, 
                    e.getMessage(), e);
            }
        });

        return result;
    }

    private String getResult(List<TextSegment> list) {
        return list.stream().map(TextSegment::text).collect(Collectors.joining("\n"));
    }
}
